import psycopg2
from django.db.utils import load_backend,DatabaseError
from .introspection import TenantSchemaDatabaseIntrospection

psql_backend=load_backend('django.db.backends.postgresql')

class DatabaseWrapper(psql_backend.DatabaseWrapper):
    def __init__(self,*args,**kwargs) -> None:
        super().__init__(*args,**kwargs)
        self.schemas=('public',)
        self.introspection=TenantSchemaDatabaseIntrospection(self)
        self.is_search_path_set=False
    @property
    def primary_schema(self):

        return self.schemas[0]
    def close(self):
        super().close()
        self.is_search_path_set=False
    def rollback(self):
        super().rollback()
        self.is_search_path_set=False
    def set_schema(self,*schemas,include_public=True):
        self.schemas=schemas
        if include_public and "public" not in schemas:
            self.schemas+=('public',)
        self.is_search_path_set=False


    def _cursor(self,name=None):
        cursor=super()._cursor(name=name)
        if self.is_search_path_set:
            return cursor
        try:
            target_path=','.join(self.schemas)
            search_path_str=f'SET search_path={target_path}'
            cursor.execute(search_path_str)
        except (DatabaseError,psycopg2.InternalError):
            self.is_search_path_set=False
        else:
            self.is_search_path_set=True
        return cursor
    
